#include <cub/cub.cuh>
#include <iostream>
typedef unsigned long long ull;
// Initialize data on the host.
#define GPU_CHECK(call)                               \
do                                                    \
{                                                     \
    const cudaError_t error_code = call;              \
    if (error_code != cudaSuccess)                    \
    {                                                 \
        printf("CUDA Error:\n");                      \
        printf("    File:       %s\n", __FILE__);     \
        printf("    Line:       %d\n", __LINE__);     \
        printf("    Error code: %d\n", error_code);   \
        printf("    Error text: %s\n",                \
            cudaGetErrorString(error_code));          \
        exit(1);                                      \
    }                                                 \
} while (0)

void initialize_data(ull *dst, unsigned int nitems) {
    // Fixed seed for illustration
    srand(2047);

    // Fill dst with random values
    for (unsigned i = 0; i < nitems; i++) {
        dst[i] = rand() % nitems;
    }
}

int main() {
    int num_segments = 7;
    int num_items = 111 * num_segments;
    // Create input data
    ull *h_data = 0;
    ull *d_data = 0;
    ull h_segment[] = {0, 111, 111 * 2, 111 * 3, 111 * 4, 111 * 5, 111 * 6, 111 * 7};
    ull *d_segment_offsets = 0;

    std::cout << "Initializing data:" << std::endl;
    h_data = (ull *)malloc(num_items * sizeof(ull));
    initialize_data(h_data, num_items);

    // 分配GPU内存
    cudaMalloc((void**)&d_data, sizeof(ull) * num_items);
    cudaMemcpy(d_data, h_data, sizeof(ull) * num_items, cudaMemcpyHostToDevice);

    cudaMalloc((void**)&d_segment_offsets, sizeof(ull) * (num_segments + 1));
    cudaMemcpy(d_segment_offsets, h_segment, sizeof(ull) * (num_segments + 1), cudaMemcpyHostToDevice);

    ull* d_sorted_data;
    GPU_CHECK(cudaMalloc((void**)&d_sorted_data, sizeof(ull) * num_items));

    // 分配排序所需的临时存储空间
    size_t temp_storage_bytes = 0;
    void* d_temp_storage = nullptr;
    cub::DeviceSegmentedRadixSort::SortKeys(d_temp_storage, temp_storage_bytes, d_data, d_sorted_data, num_items, num_segments, d_segment_offsets, d_segment_offsets + 1);

    GPU_CHECK(cudaMalloc(&d_temp_storage, temp_storage_bytes));
    
    
    // 执行排序
    std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();
    cub::DeviceSegmentedRadixSort::SortKeys(d_temp_storage, temp_storage_bytes, d_data, d_sorted_data, num_items, num_segments, d_segment_offsets, d_segment_offsets + 1);
    std::chrono::high_resolution_clock::time_point t2 = std::chrono::high_resolution_clock::now();
    std::cout<<std::chrono::duration_cast<std::chrono::microseconds>(t2-t1).count()<<std::endl;
    

    // 将排序后的结果从GPU复制回主机
    GPU_CHECK(cudaMemcpy(h_data, d_sorted_data, sizeof(ull) * num_items, cudaMemcpyDeviceToHost));

    // 打印排序后的结果
    std::cout << "Sorted data: ";
    for (int i = 0; i < 100; ++i) {
        std::cout << h_data[i] << " ";
    }
    std::cout << std::endl;

    // 释放GPU内存
    cudaFree(d_data);
    cudaFree(d_temp_storage);
    cudaFree(d_sorted_data);
}
